import datetime
from mcp.server.fastmcp import FastMCP
import re
from docx import Document
from docx.shared import Inches, Pt, RGBColor
from docx.enum.text import WD_ALIGN_PARAGRAPH
from docx.oxml.shared import OxmlElement, qn
from db.pg_connection import PostgreSQLDatabase

mcp = FastMCP("MCP Server", port=8089, host="0.0.0.0")

@mcp.tool(
    name="calculate_num",
    description="计算四则运算表达式"
)
def calculate(expression: str) -> float:  
    """计算四则运算表达式        
    参数:           
        expression: 数学表达式字符串，如 "1 + 2 * 3"        
    返回:            
    计算结果       
    """    
    try:
        # 移除表达式中的空格
        expr = expression.replace(" ", "")
        
        # 验证表达式只包含数字、小数点和基本运算符
        if not re.match(r'^[0-9+\-*/().]+$', expr):
            raise ValueError("表达式包含无效字符")
        
        # 安全地计算表达式
        result = eval(expr)
        return float(result)
    except Exception as e:
        raise ValueError(f"计算表达式 '{expression}' 时出错: {str(e)}")

@mcp.tool(
    name="get_current_time",
    description="获取当前时间"
)
def get_current_time() -> str:
    """获取当前时间        
    返回:            
    当前时间字符串       
    """    
    return datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")

# @mcp.tool(
#     name="auto_generate_word",
#     description="根据提示自动生成单词"
# )
# def auto_generate_word(prompt: str) -> str:
#     """根据提示自动生成单词        
#     参数:           
#         prompt: 提示字符串，如 "一个随机单词"        
#     返回:            
#     生成的单词字符串       
#     """    
#     # 这里简单模拟一个单词生成过程
#     return "example_word"

@mcp.tool(
    name="sql_data_extract",
    description="从SQL数据库中提取数据"
)
def sql_data_extract(sql_query: str, filename: str) -> str:
    """从SQL数据库中提取数据        
    参数:            
        sql_query: SQL查询字符串，如 "SELECT * FROM table_name"
        filename: 输出文件名（不包含扩展名）        
    返回:            
    提取的数据字符串       
    """
    try:
        # 这里简单模拟一个数据提取过程
        db = PostgreSQLDatabase()
        
        # 连接数据库
        if db.connect():
            try:
                # 查询数据
                select_query = sql_query
                results = db.execute_query(select_query)
                
                # 创建一个新的文档
                doc = Document()
                
                # 添加标题
                title = doc.add_heading(filename, 0)
                title.alignment = WD_ALIGN_PARAGRAPH.CENTER
                
                # 添加列表
                doc.add_heading('列表示例', level=1)
                doc.add_paragraph('项目1', style='List Bullet')
                doc.add_paragraph('项目2', style='List Bullet')
                doc.add_paragraph('项目3', style='List Bullet')
                
                # 添加编号列表
                doc.add_paragraph('第一步', style='List Number')
                doc.add_paragraph('第二步', style='List Number')
                doc.add_paragraph('第三步', style='List Number')
                
                # 添加表格，使用4列而不是3列
                doc.add_heading('数据表格', level=1)
                table = doc.add_table(rows=1, cols=4)
                hdr_cells = table.rows[0].cells
                hdr_cells[0].text = '区域'
                hdr_cells[1].text = '类型'
                hdr_cells[2].text = '面积'
                hdr_cells[3].text = '使用'
                
                # 添加表格数据
                if results:
                    print("查询结果:")
                    for row in results:
                        # 安全地访问元组元素
                        try:
                            # 打印结果时进行安全检查
                            row_str = f"ID: {row[0] if len(row) > 0 else 'N/A'}"
                            if len(row) > 1:
                                row_str += f", 区域: {row[1]}"
                            if len(row) > 2:
                                row_str += f", 类型: {row[2]}"
                            if len(row) > 3:
                                row_str += f", 面积: {row[3]}"
                            if len(row) > 4:
                                row_str += f", 使用: {row[4]}"
                            print(row_str)
                            
                            # 添加数据到表格时进行安全检查
                            row_cells = table.add_row().cells
                            row_cells[0].text = str(row[1]) if len(row) > 1 else ''
                            row_cells[1].text = str(row[2]) if len(row) > 2 else ''
                            row_cells[2].text = str(row[3]) if len(row) > 3 else ''
                            row_cells[3].text = str(row[4]) if len(row) > 4 else ''
                        except Exception as row_error:
                            print(f"处理行数据时出错: {str(row_error)}")
                            # 创建一个错误行
                            row_cells = table.add_row().cells
                            row_cells[0].text = "处理错误"
                            row_cells[1].text = str(row_error)
                            row_cells[2].text = ""
                            row_cells[3].text = ""
                
                # 合并表格单元格
                doc.add_heading('合并单元格示例', level=1)
                table2 = doc.add_table(rows=2, cols=2)
                table2.cell(0, 0).text = '左上'
                table2.cell(0, 1).text = '右上'
                table2.cell(1, 0).merge(table2.cell(1, 1))
                table2.cell(1, 0).text = '合并的下方单元格'
                
                # 添加分页符
                doc.add_page_break()
                
                # 添加新页面的内容
                doc.add_heading('新页面', level=1)
                doc.add_paragraph('这是新页面的内容。')
                
                # 保存文档
                doc.save(filename+".docx")
                print(f"文档已成功保存为: {filename}.docx")
                
                return f"文档已成功保存为: {filename}.docx"
            finally:
                # 断开连接
                db.disconnect()
        else:
            return "无法连接到数据库"
    except Exception as e:
        return f"执行SQL提取时出错: {str(e)}"

    return "example_data"

@mcp.tool(
    name="create_word_document",
    description="根据提示创建Word文档"
)
def create_and_edit_word_document(filename: str) -> str:
    
    # 创建一个新的文档
    doc = Document()
    
    # 添加标题
    title = doc.add_heading('Python生成的Word文档', 0)
    title.alignment = WD_ALIGN_PARAGRAPH.CENTER
    
    # 添加正文段落
    doc.add_paragraph('这是一个使用Python自动生成的Word文档示例。')
    
    # 添加带格式的段落
    paragraph = doc.add_paragraph()
    run = paragraph.add_run('这是粗体文本')
    run.bold = True
    
    run = paragraph.add_run(' 和这是普通文本，')
    
    run = paragraph.add_run('这是斜体文本')
    run.italic = True
    
    # 设置字体和颜色
    paragraph = doc.add_paragraph()
    run = paragraph.add_run('这是红色的大字体文本')
    run.font.name = '微软雅黑'
    run.font.size = Pt(16)
    run.font.color.rgb = RGBColor(255, 0, 0)
    
    # 添加列表
    doc.add_heading('列表示例', level=1)
    doc.add_paragraph('项目1', style='List Bullet')
    doc.add_paragraph('项目2', style='List Bullet')
    doc.add_paragraph('项目3', style='List Bullet')
    
    # 添加编号列表
    doc.add_paragraph('第一步', style='List Number')
    doc.add_paragraph('第二步', style='List Number')
    doc.add_paragraph('第三步', style='List Number')
    
    # 添加表格
    doc.add_heading('表格示例', level=1)
    table = doc.add_table(rows=1, cols=3)
    hdr_cells = table.rows[0].cells
    hdr_cells[0].text = '姓名'
    hdr_cells[1].text = '年龄'
    hdr_cells[2].text = '职位'
    
    # 添加表格数据
    data = [
        ('张三', '28', '工程师'),
        ('李四', '32', '产品经理'),
        ('王五', '45', '总监')
    ]
    
    for name, age, job in data:
        row_cells = table.add_row().cells
        row_cells[0].text = name
        row_cells[1].text = age
        row_cells[2].text = job
    
    # 合并表格单元格
    doc.add_heading('合并单元格示例', level=1)
    table2 = doc.add_table(rows=2, cols=2)
    table2.cell(0, 0).text = '左上'
    table2.cell(0, 1).text = '右上'
    table2.cell(1, 0).merge(table2.cell(1, 1))
    table2.cell(1, 0).text = '合并的下方单元格'
    
    # 添加分页符
    doc.add_page_break()
    
    # 添加新页面的内容
    doc.add_heading('新页面', level=1)
    doc.add_paragraph('这是新页面的内容。')
    
    # 保存文档
    doc.save(filename+".docx")
    print(f"文档已成功保存为: {filename}.docx")
    
    return f"文档已成功保存为: {filename}.docx"

if __name__ == "__main__":
    mcp.run(transport='sse')